import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, accuracy_score
from sklearn.model_selection import TimeSeriesSplit
import matplotlib.pyplot as plt

# 1. 数据获取
token = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'
ts.set_token(token)
pro = ts.pro_api()

df = pro.daily(
    ts_code='000001.SZ',
    start_date='20150101',
    end_date='20251231',
    adj='qfq'
).sort_values('trade_date')


# 2. 技术指标计算
def calculate_features(df):
    df['returns'] = df['close'].pct_change()
    df['volatility'] = df['returns'].rolling(5).std()
    df['MA5'] = df['close'].rolling(5).mean()
    df['MA10'] = df['close'].rolling(10).mean()
    df['MA20'] = df['close'].rolling(20).mean()

    delta = df['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(14).mean()
    avg_loss = loss.rolling(14).mean()
    rs = avg_gain / avg_loss
    df['RSI'] = 100 - (100 / (1 + rs))

    exp12 = df['close'].ewm(span=12, adjust=False).mean()
    exp26 = df['close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = exp12 - exp26
    df['MACD_signal'] = df['MACD'].ewm(span=9, adjust=False).mean()

    for lag in [1, 2, 3]:
        df[f'return_lag{lag}'] = df['returns'].shift(lag)

    return df.dropna()


# 3. 特征工程
df = calculate_features(df)
df['target'] = np.where(df['returns'].shift(-1) > 0, 1, 0)

features = ['volatility', 'MA5', 'MA10', 'MA20', 'RSI',
            'MACD', 'MACD_signal', 'return_lag1', 'return_lag2']
X = df[features]
y = df['target']

# 4. 初始化预测结果存储
df['prediction'] = np.nan  # 先创建全为NaN的列

# 5. 时间序列交叉验证
tscv = TimeSeriesSplit(n_splits=5)
accuracy_scores = []

for train_index, test_index in tscv.split(X):
    X_train, X_test = X.iloc[train_index], X.iloc[test_index]
    y_train, y_test = y.iloc[train_index], y.iloc[test_index]

    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=7,
        min_samples_split=10,
        class_weight='balanced',
        random_state=42
    )
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    accuracy_scores.append(accuracy_score(y_test, y_pred))

    # 将预测结果存入对应的位置
    df.iloc[test_index, df.columns.get_loc('prediction')] = y_pred

# 6. 删除没有预测结果的行
df = df.dropna(subset=['prediction'])

# 7. 计算策略收益
df['strategy_returns'] = df['returns'].shift(-1) * df['prediction']

# 8. 评估结果
print(f"平均准确率：{np.mean(accuracy_scores):.2%}")
print(classification_report(y[df.index], df['prediction']))

# 9. 可视化
cum_returns = (1 + df[['returns', 'strategy_returns']]).cumprod()

plt.figure(figsize=(12, 6))
cum_returns.plot()
plt.title('Strategy Return vs Benchmark Return')
plt.ylabel('Cumulative Return')
plt.show()


# 10. 风险指标
def calculate_metrics(returns):
    annual_return = returns.mean() * 252
    volatility = returns.std() * np.sqrt(252)
    sharpe = annual_return / volatility
    max_drawdown = (returns.cumsum().cummax() - returns.cumsum()).max()
    return pd.Series({
        '年化收益': annual_return,
        '波动率': volatility,
        '夏普比率': sharpe,
        '最大回撤': max_drawdown
    })


print("\n策略表现：")
print(calculate_metrics(df['strategy_returns'].dropna()))
print("\n基准表现：")
print(calculate_metrics(df['returns'].dropna()))